#include "Pch.h"

#include "../../MCD/Core/System/Path.h"
#include "../../MCD/Core/System/WindowEvent.h"
#include "../../MCD/Core/System/Thread.h"

#include "../../MCD/Render/Effect.h"
#include "../../MCD/Render/Material.h"
#include "../../MCD/Render/ShaderProgram.h"
#include "../../MCD/Render/RenderBufferSet.h"

#include "../Fbx/FbxFile.h"
#include "../Fbx/TriangleBufferImporter.h"

#include "../Core/RayTrace/TriangleBufferIntersect.h"
#include "../Core/RayTrace/SphericalSamplers.h"
#include "../Core/Math/Common.h"
#include "../Core/Application.h"
#include "../Core/DefaultResourceManager.h"

#include "GlVisualizer.h"

#include <vector>

using namespace IS;

namespace TestPhotonShooting
{

class Helper
{
public:
	static EffectPtr loadEffect(DefaultResourceManager& resMgr, const wchar_t* filepath)
	{
		EffectPtr eff = (Effect*)resMgr.load(filepath, true).get();

		if(nullptr != eff)
		{
			while(!eff->material.get())
				resMgr.processLoadingEvents();
		}

		return eff;
	}

	static void bindRB(RenderBufferSet& rbufs, size_t bufIdx, Effect& effect, size_t passIdx, size_t texIdx)
	{
		if(passIdx >= effect.material->getPassCount()) return;

		TextureProperty* prop = effect.material->mRenderPasses[passIdx].textureProp(texIdx);
		if(nullptr == prop) return;

		prop->texture = rbufs.bufferInfo(bufIdx).texture();
	}
};

class Lightmap
{
public:
	std::auto_ptr<RenderBufferSet> mRendBufs;

	enum RenderBufferId
	{
		E_NormalBuffer,
		E_PositionBuffer,
		E_DirectIllumBuffer,
		E_DirectIllumBuffer2,
		E_IndirectIllumBuffer,
		E_IndirectIllumBuffer2,
	};

	Lightmap(IResourceManager& resMgr, size_t size)
	{
		mRendBufs.reset(new RenderBufferSet(resMgr, size, size, RenderBufferSet::DepthBuffer_None, false));
		mRendBufs->textureBuffer(GL_RGBA32F_ARB, L"normal.rtt");
		mRendBufs->textureBuffer(GL_RGBA32F_ARB, L"position.rtt");
		mRendBufs->textureBuffer(GL_RGBA32F_ARB, L"directillum.rtt");
		mRendBufs->textureBuffer(GL_RGBA32F_ARB, L"directillum2.rtt");
		mRendBufs->textureBuffer(GL_RGBA32F_ARB, L"indirectillum.rtt");
		mRendBufs->textureBuffer(GL_RGBA32F_ARB, L"indirectillum2.rtt");
		mRendBufs->checkDeviceStatus(true);
	}
};


class TestApp : public Application
{
public:
	LinkList<TriangleBuffer> mTriBufs;
	Vec3f mLightPos;
	Vec3f mLightClr;
	bool mDrawPhotons;
	EffectPtr mGBuffEffect, mLmapEffect, mDirLtEffect, mPSplat, mShdwEffect, mDilateEffect;
	
	struct Photon
	{
		Vec3f hitPos;
		Vec3f hitNormal;
		Vec4f hitDiffuse;
		Vec4f wattage;
		Vec3f dir;
		bool active;
	};

	std::vector<Photon> mPhotons;
	std::auto_ptr<Lightmap> mLmap;
	RandomSphericalSamplerUniform mSamplerUniform;
	RandomSphericalSampler mSamplerCosine;
	
	double mGITime;
	const size_t cTargetPhotonCount;
	const size_t cPhotonPerIter;
	const size_t cMaxBounce;
	const float cDiffuseReflectivity;
	bool mDirectIllum;

	size_t mEmittedPhotonCnt;
	std::vector<Vec4f> mPhotonPos;
	std::vector<Vec4f> mPhotonDir;
	std::vector<Vec4f> mPhotonClr;

	TestApp()
		: Application(L"title=TestPhotonShooting;width=800;height=600;FSAA=1")
		, mLightPos(0, 15, 0)
		, mLightClr(1, 1, 1)
		, mDrawPhotons(false)
		//, cTargetPhotonCount(256 * 256)
		, cTargetPhotonCount(1024 * 1024)
		, cPhotonPerIter(64)
		, cMaxBounce(3)
		, cDiffuseReflectivity(1.0f)
		, mDirectIllum(true)
		, mPhotonPos(cPhotonPerIter)
		, mPhotonDir(cPhotonPerIter)
		, mPhotonClr(cPhotonPerIter)
	{
		mCamera.moveForward(-10.0f);

		// load effects
		mGBuffEffect = Helper::loadEffect(mResMgr, L"gbuffersetup.fx.xml");
		mLmapEffect = Helper::loadEffect(mResMgr, L"lightmap.fx.xml");
		mDirLtEffect = Helper::loadEffect(mResMgr, L"directlighting.fx.xml");
		mPSplat = Helper::loadEffect(mResMgr, L"photonsplatting.fx.xml");
		mShdwEffect = Helper::loadEffect(mResMgr, L"shadowmapsetup.fx.xml");
		mDilateEffect = Helper::loadEffect(mResMgr, L"dilation.fx.xml");

		// open fbx file
		std::wstring fbxfilepath = (Path::getCurrentPath() / L"media/BoxRoom.fbx").getString();

		Log::format( Log::Info, L"Importing %s ...", fbxfilepath.c_str() );

		FbxFile fbxfile(fbxfilepath.c_str());

		if(!fbxfile.open())
			return;

		Log::format( Log::Info, L"FBX file version %d.%d.%d"
			, fbxfile.fileVersion().major
			, fbxfile.fileVersion().minor
			, fbxfile.fileVersion().revision
			);

		// import triangle buffers
		TriangleBufferImporter().import(fbxfile, mTriBufs);

		// create light map
		mLmap.reset(new Lightmap(mResMgr, 512));

		gi_Reset();
	}

	void drawPhotons()
	{
		if(!mDrawPhotons) return;

		glPushAttrib(GL_LIGHTING_BIT|GL_DEPTH_BUFFER_BIT);
		//glDisable(GL_DEPTH_TEST);
		glDisable(GL_LIGHTING);

		glBegin(GL_LINES);

		const float s = 0.1f;
		
		const size_t cPhotonCnt = mPhotons.size();
		for(size_t i = 0; i < cPhotonCnt; ++i)
		{
			Photon* p = &mPhotons[i];
			glColor3f(p->wattage.x, p->wattage.y, p->wattage.z);
			GlVisualizer::drawPoint(p->hitPos, s);
			//GlVisualizer::drawLine(mLightPos, p->hitPos);
		}

		glColor3f(0.0f, 1.0f, 0.0f);
		GlVisualizer::drawPoint(mLightPos, s * 5.0f);

		glEnd();

		glPopAttrib();
	}

	void drawScene(Effect& effect)
	{
		for(TriangleBuffer* tbuf = mTriBufs.begin(); tbuf != mTriBufs.end(); tbuf = tbuf->next())
		{
			glPushMatrix();
			glMultTransposeMatrixf(tbuf->localTransform.getPtr());

			if(&effect == mGBuffEffect.get())
			{
				effect.material->preRender(0);
				ShaderProgram::current()->uniformMatrix4fv("g_WorldMatrix", 1, true, tbuf->localTransform.getPtr());
				ShaderProgram::current()->uniform1f("g_TextureSpace", 1.0f);
				tbuf->draw();
				effect.material->postRender(0);
			}
			else if(&effect == mShdwEffect.get())
			{
				const float clipN = this->mCamera.frustum.near;
				const float clipF = this->mCamera.frustum.far;

				effect.material->preRender(0);
				ShaderProgram::current()->uniform4f("g_ClipPlanes", clipN, clipF, clipF - clipN, 1 / (clipF - clipN));
				tbuf->draw();
				effect.material->postRender(0);
			}
			else if(&effect == mLmapEffect.get())
			{
				effect.material->preRender(0);
				if(0 == mEmittedPhotonCnt)
					ShaderProgram::current()->uniform1f("g_indirectIllumScale", 0.0f);
				else
					ShaderProgram::current()->uniform1f("g_indirectIllumScale", (float)(1 / (double)mEmittedPhotonCnt));

				ShaderProgram::current()->uniform1f("g_directIllumScale", 1.0f);

				tbuf->draw();
				effect.material->postRender(0);
			}
			else
			{
				effect.material->preRender(0);
				tbuf->draw();
				effect.material->postRender(0);
			}

			glPopMatrix();
		}
	}

	void drawSceneToScreen()
	{
		glClearColor(0.3f, 0.3f, 0.3f, 0.0f);
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

		Helper::bindRB(*mLmap->mRendBufs, Lightmap::E_DirectIllumBuffer2, *mLmapEffect, 0, 0);
		Helper::bindRB(*mLmap->mRendBufs, Lightmap::E_IndirectIllumBuffer2, *mLmapEffect, 0, 1);
		drawScene(*mLmapEffect);
	}
	
	void drawSceneToShadowmap()
	{
		glClearColor(0.0f, 0.0f, 0.0f, 0.0f);
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

		Helper::bindRB(*mLmap->mRendBufs, Lightmap::E_DirectIllumBuffer, *mShdwEffect, 0, 0);
		drawScene(*mShdwEffect);
	}

	void gi_Reset()
	{
		mPhotons.clear();
		if(mDrawPhotons)
			mPhotons.reserve( cTargetPhotonCount * (cMaxBounce + 1) );
		mEmittedPhotonCnt = 0;
		mGITime = 0;

		gi_InitGBuffer();
		gi_ProcessDirectLighting();
	}
	
	void gi_InitGBuffer()
	{
		mLmap->mRendBufs->begin(Lightmap::E_NormalBuffer, Lightmap::E_PositionBuffer);
		mLmap->mRendBufs->checkDeviceStatus(false);

		glClearColor(0.0f, 0.0f, 0.0f, 0.0f);
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

		drawScene(*mGBuffEffect);

		mLmap->mRendBufs->end();
	}

	void gi_ProcessDirectLighting()
	{
		Helper::bindRB(*mLmap->mRendBufs, Lightmap::E_NormalBuffer, *mDirLtEffect, 0, 0);
		Helper::bindRB(*mLmap->mRendBufs, Lightmap::E_PositionBuffer, *mDirLtEffect, 0, 1);

		mLmap->mRendBufs->begin(Lightmap::E_DirectIllumBuffer);
		mLmap->mRendBufs->checkDeviceStatus(false);

		glClearColor(0.0f, 0.0f, 0.0f, 0.0f);
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

		mDirLtEffect->material->preRender(0);

		ShaderProgram::current()->uniform3fv("g_lightPos", 1, mLightPos.data);
		ScreenQuad::draw(mLmap->mRendBufs->target(), 0, 0, mLmap->mRendBufs->width(), mLmap->mRendBufs->height());
		mDirLtEffect->material->postRender(0);

		mLmap->mRendBufs->end();
	}

	void gi_ShootPhotons()
	{
		if(mEmittedPhotonCnt >= cTargetPhotonCount)
			return;

		if(mEmittedPhotonCnt == 0)
			Log::format(Log::Info, L"Start Photon shooting cTargetPhotonCount=%d; cMaxBounce=%d; mDirectIllum=%s..."
				, cTargetPhotonCount
				, cMaxBounce
				, mDirectIllum ? L"true" : L"false"
				);
		
		DeltaTimer t;
		
		// gpu init
		Helper::bindRB(*mLmap->mRendBufs, Lightmap::E_NormalBuffer, *mPSplat, 0, 0);
		Helper::bindRB(*mLmap->mRendBufs, Lightmap::E_PositionBuffer, *mPSplat, 0, 1);

		mLmap->mRendBufs->begin(Lightmap::E_IndirectIllumBuffer);
		mLmap->mRendBufs->checkDeviceStatus(false);

		if(mEmittedPhotonCnt == 0)
		{
			glClearColor(0.0f, 0.0f, 0.0f, 0.0f);
			glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);
		}

		mPSplat->material->preRender(0);

		// cpu init
		std::vector<Photon> photons(cPhotonPerIter);

		std::vector<Vec3f> samples; samples.reserve(cPhotonPerIter);
		mSamplerUniform.generateN(cPhotonPerIter, samples);

		TriangleBufferIntersect tbi;
		tbi.beginIntersect(mTriBufs);

		const float cOffset = 0.1f;
		const float cRandNormScale = 0.25f;

		// emit photons from light source
		for(size_t i=0; i<cPhotonPerIter; ++i)
		{
			Photon& p = photons[i];
			
			size_t hitCnt = tbi.intersect(mLightPos, samples[i], false);

			if(0 == hitCnt)
			{
				p.active = false;
				continue;
			}

			const size_t cHitIdx = tbi.getClosestHitIndex();

			p.active = true;
			p.dir = samples[i];
			p.hitNormal = tbi.getHitNormal(cHitIdx);
			p.hitPos = tbi.getHitPosition(cHitIdx) + p.hitNormal * cOffset;
			p.wattage = Vec4f(mLightClr.x, mLightClr.y, mLightClr.z, 1);
			p.hitDiffuse = tbi.getHitColor(cHitIdx);
			p.hitDiffuse.x *= p.hitDiffuse.x;
			p.hitDiffuse.y *= p.hitDiffuse.y;
			p.hitDiffuse.z *= p.hitDiffuse.z;
		}
		
		mEmittedPhotonCnt += cPhotonPerIter;

		if(mDirectIllum)
		{
			samples.clear();
			mSamplerUniform.generateN(cPhotonPerIter, samples);

			for(size_t i=0; i<cPhotonPerIter; ++i)
			{
				// some init value
				mPhotonClr[i] = Vec4f(0, 0, 0, 0);

				Photon& p = photons[i];
				if(!p.active) continue;

				if(samples[i].dot(p.hitNormal) < 0)
					samples[i] = -samples[i];

				samples[i] = (cRandNormScale * p.hitNormal + samples[i]).normalizedCopy();
				p.dir = samples[i];

				mPhotonPos[i] = Vec4f(p.hitPos, 0);
				mPhotonDir[i] = Vec4f(p.dir, 0);
				mPhotonClr[i] = p.wattage;

				if(mDrawPhotons) mPhotons.push_back(p);
			}

			ShaderProgram::current()->uniform4fv("g_photonPosition", cPhotonPerIter, (float*)(mPhotonPos[0].data));
			ShaderProgram::current()->uniform4fv("g_photonDirection", cPhotonPerIter, (float*)mPhotonDir[0].data);
			ShaderProgram::current()->uniform4fv("g_photonColor", cPhotonPerIter, (float*)(mPhotonClr[0].data));
			
			ScreenQuad::draw(mLmap->mRendBufs->target(), 0, 0, mLmap->mRendBufs->width(), mLmap->mRendBufs->height());
		}

		// photon bouncing
		for(size_t ibounce=0; ibounce<cMaxBounce; ++ibounce)
		{
			samples.clear();
			mSamplerUniform.generateN(cPhotonPerIter, samples);

			for(size_t i=0; i<cPhotonPerIter; ++i)
			{
				// some init value
				mPhotonClr[i] = Vec4f(0, 0, 0, 0);

				Photon& p = photons[i];
				if(!p.active) continue;

				size_t hitCnt = tbi.intersect(p.hitPos, p.dir, false);
				if(0 == hitCnt)
				{
					p.active = false;
					continue;
				}
				
				const size_t cHitIdx = tbi.getClosestHitIndex();
				p.hitNormal = tbi.getHitNormal(cHitIdx);
				p.hitPos = tbi.getHitPosition(cHitIdx) + p.hitNormal * cOffset;

				if(samples[i].dot(p.hitNormal) < 0)
					samples[i] = -samples[i];

				samples[i] = (cRandNormScale * p.hitNormal + samples[i]).normalizedCopy();
				p.dir = samples[i];

				p.wattage.x *= p.hitDiffuse.x * cDiffuseReflectivity;
				p.wattage.y *= p.hitDiffuse.y * cDiffuseReflectivity;
				p.wattage.z *= p.hitDiffuse.z * cDiffuseReflectivity;

				p.hitDiffuse = tbi.getHitColor(cHitIdx);
				p.hitDiffuse.x *= p.hitDiffuse.x;
				p.hitDiffuse.y *= p.hitDiffuse.y;
				p.hitDiffuse.z *= p.hitDiffuse.z;

				mPhotonPos[i] = Vec4f(p.hitPos, 0);
				mPhotonDir[i] = Vec4f(p.dir, 0);
				mPhotonClr[i] = p.wattage;

				if(mDrawPhotons) mPhotons.push_back(p);
			}

			{
				ShaderProgram::current()->uniform4fv("g_photonPosition", cPhotonPerIter, (float*)(mPhotonPos[0].data));
				ShaderProgram::current()->uniform4fv("g_photonDirection", cPhotonPerIter, (float*)mPhotonDir[0].data);
				ShaderProgram::current()->uniform4fv("g_photonColor", cPhotonPerIter, (float*)(mPhotonClr[0].data));
				
				ScreenQuad::draw(mLmap->mRendBufs->target(), 0, 0, mLmap->mRendBufs->width(), mLmap->mRendBufs->height());
			}
		}

		// cpu clean up
		tbi.endIntersect();

		// gpu clean up
		mPSplat->material->postRender(0);
		mLmap->mRendBufs->end();

		// timing
		mGITime += t.getDelta().asSecond();

		mSleep(1);

		if(mEmittedPhotonCnt >= cTargetPhotonCount)
			Log::format(Log::Info
				, L"All %d Photons have been shot in %f seconds."
				, cTargetPhotonCount
				, mGITime);
		else
			wprintf(L"Photons: %d/%d; FPS:%3.2f\r", mEmittedPhotonCnt, cTargetPhotonCount, mLastFPS);
	}

	void gi_Dilation(Lightmap::RenderBufferId srcBufId, Lightmap::RenderBufferId dstBufId)
	{
		Helper::bindRB(*mLmap->mRendBufs, srcBufId, *mDilateEffect, 0, 0);

		mLmap->mRendBufs->begin(dstBufId);
		mLmap->mRendBufs->checkDeviceStatus(false);

		glClearColor(0.0f, 0.0f, 0.0f, 0.0f);
		glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

		mDilateEffect->material->preRender(0);

		ShaderProgram::current()->uniform4f("g_texelSize", 1.0f / mLmap->mRendBufs->width(), 1.0f / mLmap->mRendBufs->height(), 0, 0 );
		ScreenQuad::draw(mLmap->mRendBufs->target(), 0, 0, mLmap->mRendBufs->width(), mLmap->mRendBufs->height());
		mDilateEffect->material->postRender(0);

		mLmap->mRendBufs->end();
	}

	sal_override void update(float deltaTime)
	{
		gi_ShootPhotons();

		gi_Dilation(Lightmap::E_DirectIllumBuffer, Lightmap::E_DirectIllumBuffer2);
		gi_Dilation(Lightmap::E_IndirectIllumBuffer, Lightmap::E_IndirectIllumBuffer2);

		drawSceneToScreen();
		//drawSceneToShadowmap();

		drawPhotons();
	}

	sal_override void onEvent(const MCD::Event& e)
	{
		Application::onEvent(e);

		if(Event::KeyPressed == e.Type && Key::H == e.Key.Code)
		{
			wprintf(L"F1-restart\nF2-display photons\nF3-apply direct illumination photons\n");
		}

		if(Event::KeyPressed == e.Type && Key::F1 == e.Key.Code)
			gi_Reset();

		if(Event::KeyPressed == e.Type && Key::F2 == e.Key.Code)
		{
			mDrawPhotons = !mDrawPhotons;
			//wprintf(L"mDrawPhotons = %s\n", (mDrawPhotons?L"true":L"false") );
		}
		
		if(Event::KeyPressed == e.Type && Key::F3 == e.Key.Code)
		{
			mDirectIllum = !mDirectIllum;
			wprintf(L"mDirectIllum = %s\n", (mDirectIllum?L"true":L"false") );
		}
	}

};	// class TestApp

}	// namespace TestPhotonShooting

TEST(TestPhotonShooting)
{
	TestPhotonShooting::TestApp app;
	app.mainLoop();
}